package main

import "container/heap"

// 无向加权图
// edge[i]=[a,b]    succProb[i]

// Dijkstra 单源最短  贪心

func maxProbability(n int, edges [][]int, succProb []float64, start int, end int) float64 {
	graph := make([][]state, n)
	for i := 0; i < len(edges); i++ {
		edge := edges[i]
		graph[edge[0]] = append(graph[edge[0]], state{edge[1], succProb[i]})
		graph[edge[1]] = append(graph[edge[1]], state{edge[0], succProb[i]})
	}
	dist := make([]float64, n)
	dist[start] = 1
	h := &hp{state{start, 1}}
	for 0 < h.Len() {
		cur := heap.Pop(h).(state)
		if cur.id == end {
			return cur.dist
		}
		if cur.dist < dist[cur.id] { // 当前路径小于最大距离，舍弃
			continue
		}
		for _, next := range graph[cur.id] {
			//如果从start到cur节点再到next节点的距离大于已知start到next最大距离，则更新
			if d := cur.dist * next.dist; d > dist[next.id] {
				dist[next.id] = d
				heap.Push(h, state{next.id, d})
			}
		}
	}
	return dist[end]
}

type state struct {
	id   int
	dist float64
}

type hp []state

func (h hp) Len() int            { return len(h) }
func (h hp) Less(i, j int) bool  { return h[i].dist > h[j].dist }
func (h hp) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *hp) Push(v interface{}) { *h = append(*h, v.(state)) }
func (h *hp) Pop() (v interface{}) {
	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
	return
}
